import requests
import openai_auth
import json
import pickle
import random
from tqdm import tqdm
import time
import copy
import os
from utils import drug_names, emerging_euph


# ----------------------------------step 1--------------------------------------------
# 确定好每个毒品的委婉语是什么


def step1():
    output = dict()
    drugnames = set(list(json.load(open("drugeuphs.json", "r", encoding="utf-8")).keys()))
    for d in drugnames:

        data = {
            "model": "gpt-3.5-turbo",
            "temperature": 0.7,
            "top_p": 1.0,
            "frequency_penalty": 0.0,
            "presence_penalty": 0.0,
            "max_tokens": 1500,
            "messages": [{"role": "system",
                          "content": """We are conducting research on drug euphemisms. The user will input a drug name. The assistant should generate 10 Euphemisms for that drug. Drug euphemisms should not be related to any drug or chemical acronyms. The assistant should give a euphemism list without any other sentences."""},
                         {"role": "user", "content": d}, ]
        }
        HEADER = {"Authorization": f"Bearer {openai_auth.OpenAiKey}"}
        response = requests.post(f"https://api.openai.com/v1/chat/completions", json=data, headers=HEADER,
                                 proxies={"https": "127.0.0.1:7890"})
        if "error" in response.json():
            return response.json()
        _ = response.json()["choices"][0]["message"]["content"].split("\n")
        _ = set([each.split(".")[-1].strip().lower().replace(" ", "_") for each in _])
        for i in list(_):
            if len(i) == 0:
                _.discard(i)
            for w in ["pill", "drug", "weed", "xan", "oxy", "heroin", "ecstasy", "meth"]:
                if w in i:
                    _.discard(i)
        _ = list(_ - set(drug_names) - set(emerging_euph.keys()))  # 去除已知的
        output[d] = _
        print(f"{d}:{_}")
        # pickle.dump(_, open(f"./{d}.pkl", "wb"))
        time.sleep(30)
    return output


drug2euph_dic = step1()
pickle.dump(drug2euph_dic, open("./drug2euph_dic.pkl", "wb"))

# # ---------------------------------------step 2---------------------------------------------------------------------------
request_body = {
    "model": "gpt-3.5-turbo",
    "temperature": 0.7,
    "top_p": 0.6,
    "frequency_penalty": 0.7,
    "presence_penalty": 0.3,
    "max_tokens": 1000,
    "messages": [{"role": "system",
                  # "content": """We are conducting research on drug euphemisms. The user will input a drug sentence with a specified drug name. The assistant should rewrite the sentence, using a drug euphemism to moderate the drug name. After that, the assistant should generate another harmless sentence containing the exact drug euphemism used in the rewritten sentence but has its original normal meaning. The harmless sentence should look like drug sentence."""},
                  "content": """We are conducting research on drug euphemisms. The user will input a drug sentence with a specified drug name and the drug's euphemism. The assistant should do 2 tasks: (1) rewrite the sentence, using the given drug euphemism to moderate the drug name; (2) generate another harmless sentence containing the exact drug euphemism used in the rewritten sentence but has its original normal meaning.

The harmless sentence should look like the given drug sentence but do not contain any drug information like drug tabs. At least the harmless sentence should have almost the same length as the drug sentence.

Remember the euphemism must be used both in the rewritten sentence and the harmless sentence. """},
                 {"role": "user",
                  "content": """Drug name: amphetamine

Drug sentence: way to fucking much man dont_ask for help when its clear the reason_why man you_should have_known that wasnt good for you mdma has enough amphetamine in it why_would you_wanna take_more?

Euphemism: speed"""},
                 {"role": "assistant",
                  "content": """Rewritten sentence: Way too much man! Don't ask for help when it's clear the reason why. You should have known that wasn't good for you. Molly has enough {speed} in it, why would you wanna take more?

Harmless sentence: Way too much man! Don't ask for help when it's clear the reason why. Its {speed} is at its limit, which makes it very difficult to maneuver, and of course, accidents will happen."""},

                 #                  {"role": "user",
                 #                   "content": """Drug name: opium
                 #
                 # Drug sentence: i want know_what the high feels_like but i also never want_to_touch it ever my compromise is never trying it but someday trying opium_tea? smoked? to experience a taste of what it's like without it being so overblown that i desire it_again i'd_be fine trying opium once "well that was fun " check_box don't_think i'd_be fine trying heroin once because if i go to the trouble_of_finding a way to get that_shit it_won't be once only_problem is where in the fuck do you get opium in the states? it's_not exactly a common drug here guess_i'll have to give one of the new tor_markets a shot."""},
                 #                  {"role": "assistant",
                 #                   "content": """Euphemism: tar
                 #
                 # Rewritten sentence: I want to know what the high feels like, but I also never want to touch it ever. My compromise is never trying it but someday trying {tar} tea? Smoked? To experience a taste of what it's like without it being so overblown that I desire it again. I'd be fine trying {tar} once, 'well that was fun' check box. Don't think I'd be fine trying heroin once because if I go to the trouble of finding a way to get that stuff, it won't be once. Only problem is where in the fuck do you get tar in the states? It's not exactly a common drug here. Guess I'll have to give one of the new tor markets a shot.
                 #
                 # Harmless sentence: I want to know what cooking with {tar} feels like, but I also never want to touch it ever. My compromise is never trying it but someday trying {tar} tea? To experience a taste of what it's like without it being so overblown that I desire it again. I'd be fine trying {tar} tea once, 'well that was fun' check box. Don't think I'd be fine trying tea once because if I go to the trouble of finding a way to get that stuff, it won't be once. Only problem is where in the fuck do you get {tar} in the states? It's not exactly a common  here. Guess I'll have to give one of the new tor markets a shot."""},

                 {"role": "user",
                  "content": """Drug name: marijuana

Drug sentence: skimmed but i_wasn't referring to marijuana at all actually_believe it or not i_am pretty_well versed_in_pharmacology amateur and do_not feel_the_need to defend anything also i_am not a teenager the worst effects of adderall that i_have heard of are on long_term memory and bone and heart strength iirc my 15mg doesn't_hurt me that said i_am not_a_fan of amphetamines and don't_enjoy it as_much at all really as i do derive other benefit from it sorry i_am not up to matching your_argument length/depth today.

Euphemism: weed"""},
                 {"role": "assistant",
                  "content": """Rewritten sentence: Skimmed, but I wasn't referring to {weed} at all. Actually, believe it or not, I am pretty well-versed in pharmacology as an amateur and do not feel the need to defend anything. Also, I am not a teenager. The worst effects of Adderall that I have heard of are on long-term memory and bone and heart strength, iirc. My 15mg doesn't hurt me. That said, I am not a fan of amphetamines and don't enjoy it as much at all really as I do derive other benefit from it. Sorry, I'm not up to matching your argument length/depth today.

Harmless sentence: Skimmed, but I wasn't referring to {weed} at all. Actually, believe it or not, I am pretty well-versed in botany as an amateur and do not feel the need to defend anything. Also, I am not a teenager. The worst effect of gardening is that it can be very time-consuming. Gardening requires a lot of maintenance, such as weeding, pruning, planting, and watering. If you don't keep up with it, it can quickly become overgrown and unkempt. Sorry, I'm not up to matching your argument length/depth today."""},

                 {"role": "user",
                  "content": """Drug name: heroin

Drug sentence: it_was his_choice to start opiates but_unfortunately this_story is super_common my_friend who_claimed would_never do heroin got_introduced to it in court_appointed rehab because it_gets out of your_system quick he probably would_have been perfectly_content with just smoking up to that point yeah it could_have happened without rehab but you_are a product of your_environment you_have to take into_consideration the influences people have on you sure he_could have_taken the high road but this_is hard even for strong_willed people.

Euphemism: smack"""},
                 {"role": "assistant",
                  "content": """Rewritten sentence: It was his choice to start opiates, but unfortunately, this story is super common. My friend who claimed he would never do {smack} got introduced to it in court-appointed rehab because it gets out of your system quick. He probably would have been perfectly content with just smoking up to that point. Yeah, it could have happened without rehab, but you are a product of your environment. You have to take into consideration the influences people have on you. Sure, he could have taken the high road, but this is hard even for strong-willed people.

Harmless sentence: It was his choice to start exercising, but unfortunately, this story is super common. My friend who claimed he would never do {smack} got introduced to it in a fitness center because it's an effective way of losing weight quickly. He probably would have been perfectly content with just lifting weights up to that point. Yeah, it could have happened without the gym membership, but you are a product of your environment. You have to take into consideration the influences people have on you. Sure, he could have taken the high road and started jogging or cycling instead, but this is hard even for strong-willed people who hate cardio workouts like running or cycling."""},

                 {"role": "user",
                  "content": """Drug name: ecstasy

Drug sentence: yeah_i haven\'t_watched it in a while but you're_probably right he_may not have named the drug empathy but definitely does characterize_the_effects of mdma as empathy instead of ecstasy but society twisted that "good" name into ecstasy which shulgin_didn\'t appreciate.

Euphemism: beans"""},
                 {"role": "assistant",
                  "content": """Rewritten sentence: Yeah, I haven't watched it in a while, but you're probably right. He may not have named the drug empathy, but he definitely does characterize the effects of MDMA as empathy instead of {beans}. But society twisted that "good" name into {beans}, which Shulgin didn't appreciate.

Harmless sentence: Yeah, I haven't eaten it in a while, but you're probably right. He may not have named the dish {beans}, but he definitely does characterize its flavor as spicy instead of bland. But society twisted that "good" taste into something else entirely different from what it really is."""},
                 ]
}

request_body_2 = {
    "model": "gpt-3.5-turbo",
    "temperature": 0.7,
    "top_p": 0.6,
    "frequency_penalty": 0.7,
    "presence_penalty": 0.3,
    "max_tokens": 1000,
    "messages": [{"role": "system",
                  # "content": """We are conducting research on drug euphemisms. The user will input a drug sentence with a specified drug name. The assistant should rewrite the sentence, using a drug euphemism to moderate the drug name. After that, the assistant should generate another harmless sentence containing the exact drug euphemism used in the rewritten sentence but has its original normal meaning. The harmless sentence should look like drug sentence."""},
                  "content": """We are conducting research on drug euphemisms. The user will input a drug sentence with a specified drug name and the drug's euphemism. The assistant should do 2 tasks: (1) rewrite the sentence, using the given drug euphemism to moderate the drug name; (2) generate another harmless sentence containing the exact drug euphemism used in the rewritten sentence but has its original normal meaning.

The harmless sentence should look like the given drug sentence but do not contain any drug information like drug tabs. At least the harmless sentence should have almost the same length as the drug sentence.

Remember the euphemism must be used both in the rewritten sentence and the harmless sentence. """},
                 {"role": "user",
                  "content": """Word: lucy

Sentence: Dude, sounds like a perfect plan. You yourself said you have taken 2x the amount of sugar in your coffee at a cafe and were fine. Half that amount with some creamer sounds like you'll be able to handle yourself fine. Just bring water, gum, and Vicks for fresh breath and comfort. Get there early if the place is crowded; I've been there before, and it's pretty busy during peak hours."""},
                 {"role": "assistant",
                  "content": """Rewritten Sentence: Dude, {Lucy}, sounds like a perfect plan. You yourself said you have taken 2x the amount of sugar in your coffee at a cafe and were fine. Half that amount with some creamer sounds like you'll be able to handle yourself fine. Just bring water, gum, and Vicks for fresh breath and comfort. Get there early if the place is crowded; I've been there before, and it's pretty busy during peak hours."""},
                 {"role": "user",
                  "content": """Word: acid

Sentence: That part is untrue. In cool, dry, and dark conditions, vinegar will last years. That guy who studies vinegar who did an AMA awhile back said he's got some that's years old that hasn't lost any acidity. I can't find his name or I'd find his comment. If this is fact then I stand corrected about the shelf life of vinegar in ideal storage conditions."""},
                 {"role": "assistant",
                  "content": """Rewritten Sentence: That part is untrue. In cool, dry, and dark conditions, {acid} like vinegar will last years. That guy who studies vinegar who did an AMA awhile back said he's got some that's years old that hasn't lost any acidity. I can't find his name or I'd find his comment. If this is fact then I stand corrected about the shelf life of vinegar in ideal storage conditions."""},
                 ]
}


def chatGPTrequest(inputs, rewriteagain=False, data=None):
    """

    :param inputs: [Drug name, drugsentence] rewriteagain=False || [Word, Sentence] rewriteagain=True
    :param data:
    :return:
    """
    if not rewriteagain:
        if data is None:
            data = copy.deepcopy(request_body)
        data["messages"].append({
            "role": "user",
            "content": f"""Drug name: {inputs[0]}\nDrug sentence:{inputs[1]}\nEuphemism: {inputs[2]}"""
        })
    else:
        if data is None:
            data = copy.deepcopy(request_body_2)
        data["messages"].append({
            "role": "user",
            "content": f"""Word: {inputs[0]}\nSentence:{inputs[1]}"""
        })
    HEADER = {"Authorization": f"Bearer {openai_auth.OpenAiKey}"}
    response = requests.post(f"https://api.openai.com/v1/chat/completions", json=data, headers=HEADER,
                             proxies={"https": "127.0.0.1:7890"})
    if "error" in response.json():
        return response.json()
    output = response.json()["choices"][0]["message"]
    return output


def get_devset():
    drugnames = set(list(json.load(open("drugeuphs.json", "r", encoding="utf-8")).keys()))
    # drugnames = drugnames - {"alprazolam", "amphetamine", "ecstasy", "hydrocodone", "ketamine", "marijuana",
    #                          "marijuana_concentrates", "mescaline", "opium", "oxycodone", "pcp", "promethazine",
    #                          "psilocybin_mushrooms", "steroids",
    #                          "synthetic_cathinones", }

    black = pickle.load(open("./corpus/black_corp.pkl", "rb"))
    black_dic = pickle.load(open("./corpus/black_dic.pkl", "rb"))
    drug2euph_dic = pickle.load(open("./drug2euph_dic.pkl", "rb"))

    for drug in tqdm(drugnames):
        output = []  # [[Drug, Drug Sentence, Euphemism, Rewritten Sentence, Harmless Sentence]]
        _ = []
        while len(output) != 3:
            choice = random.choice(black_dic[drug])
            while len(black[choice]) > 300 or len(black[choice]) < 20:
                choice = random.choice(black_dic[drug])

            euphemism = random.choice(drug2euph_dic[drug])
            drugsentence = " ".join(black[choice])
            _ = [drug, drugsentence, euphemism]
            time.sleep(40)
            print([drug, drugsentence, euphemism])
            try:
                res = chatGPTrequest([drug, drugsentence, euphemism])
            except:
                time.sleep(10)
                res = chatGPTrequest([drug, drugsentence, euphemism])
            res["content"] = res["content"].lower()
            print(res)
            res = res["content"].split("\n")
            for line in res:
                if ":" in line:
                    _.append(line.split(":", 1)[1].strip())
            if len(_) != 5:
                continue
            cnt, cnt_, flag = 0, 3, True  # 如果harmless句子没有出现委婉语，就需要重新生成，且只有cnt_次机会。
            while "{" + euphemism.replace("_", " ") + "}" not in _[4]:
                time.sleep(40)
                try:
                    res = chatGPTrequest([_[2], _[4]], rewriteagain=True)
                except:
                    time.sleep(10)
                    res = chatGPTrequest([_[2], _[4]], rewriteagain=True)
                _[4] = res["content"].lower().split("rewritten sentence:")[-1].strip()
                cnt += 1
                if cnt >= cnt_:
                    flag = False
                    break
            if flag:
                output.append(_)
            else:
                continue

        pickle.dump(output, open(f"./drug_{drug}.pkl", "wb"))


get_devset()
output = []
for each in os.listdir():
    if "drug_" in each:
        _ = pickle.load(open(each, "rb"))
        output += _
pickle.dump(output, open("./corpus/devset_org.pkl", "wb"))

# ---------------------------------------step 3---------------------------------------------------------------------------
# 再次检查gpt生成是否准确，如果有错误就再改
dev = pickle.load(open("./corpus/devset_org.pkl", "rb"))
for idx in range(len(dev)):
    if "harmless sentence:" in dev[idx][4]:
        while True:
            s = dev[idx][4].split("harmless sentence:")
            if "{" + dev[idx][2] + "}" in s[0] or "{" + dev[idx][2].replace("_", " ") + "}" in s[0]:
                dev[idx][4] = s[0]
                break
            else:
                while True:
                    try:
                        time.sleep(40)
                        res = chatGPTrequest([dev[idx][2], s[0]], rewriteagain=True)
                    except:
                        continue
                    if "content" in res.keys():
                        break
                    else:
                        continue
                dev[idx][4] = res["content"].lower().split("rewritten sentence:")[1].strip()
                if "harmless sentence:" in dev[idx][4]:
                    print(idx, dev[idx][2], dev[idx][4])
                    continue
                else:
                    break

pickle.dump(dev, open("./corpus/devset_org.pkl", "wb"))

dev = pickle.load(open("corpus/devset.pkl", "rb"))
for idx in range(len(dev)):
    dev[idx][1] = dev[idx][1].replace(dev[idx][0], "{" + dev[idx][0] + "}")
    dev[idx][4] = dev[idx][4].replace(dev[idx][2].replace("_", " "), dev[idx][2])
    dev[idx][3] = dev[idx][3].replace(dev[idx][2].replace("_", " "), dev[idx][2])
pickle.dump(dev, open("./corpus/devset.pkl", "wb"))

dev = pickle.load(open("corpus/devset.pkl", "rb"))
print("pass")
